import base
import imageio
import os
import glob
import numpy as np
import math


def get_focused_area_list(img, gt):
    img_list = []
    gt_list = []

    gt_h, gt_w, gt_c = gt.shape

    gt_size = math.ceil(max(int(gt_w / 3), int(gt_h * 2 / 3))/16)*16

    step = gt_size // 4
    for offset_gt_h in range(0, gt_h - gt_size, step):
        for offset_gt_w in range(0, gt_w - gt_size, step):
            one_img = img[offset_gt_h:offset_gt_h + gt_size,
                          offset_gt_w:offset_gt_w + gt_size, :]
            one_gt = gt[offset_gt_h:offset_gt_h + gt_size,
                        offset_gt_w:offset_gt_w + gt_size, :]
            img_list.append(one_img)
            gt_list.append(one_gt)
    return img_list, gt_list


def get_image(vi_name, gt_name, crop=0):
    vi = imageio.imread(vi_name)
    gt = imageio.imread(gt_name)
    assert np.shape(vi) == np.shape(gt), f"{vi_name}-{gt_name}大小不相同"
    h, w, c = np.shape(vi)
    if (crop != 0):
        vi = vi[crop:h - crop, crop:w - crop, :]
        gt = gt[crop:h - crop, crop:w - crop, :]
    return vi, gt


def get_focused_psnr(vi_path, gt_path, mode="RGB", crop=0):
    vi_list = sorted(glob.glob(os.path.join(vi_path, "*.png"))+glob.glob(os.path.join(vi_path, "*.jpg")))
    gt_list = sorted(glob.glob(os.path.join(gt_path, "*.png"))+glob.glob(os.path.join(gt_path, "*.jpg")))
    if (len(vi_list) != len(gt_list)):
        raise ValueError("视频长度不匹配")
    total_psnr = 0
    for vi_name, gt_name in zip(vi_list, gt_list):
        vi, gt = get_image(vi_name, gt_name, crop=crop)
        img_list, gt_list = get_focused_area_list(vi, gt)
        psnr_list = []
        for one_img, one_gt in zip(img_list, gt_list):
            psnr_list.append(base.psnr(one_img, one_gt, mode=mode))
        assert psnr_list != [], 'psnr_list 不能为空'
        total_psnr += np.mean(psnr_list)

    return total_psnr / len(vi_list)


def get_psnr(vi_path, gt_path, mode="RGB", crop=0):
    vi_list = sorted(glob.glob(os.path.join(vi_path, "*.png")))
    gt_list = sorted(glob.glob(os.path.join(gt_path, "*.png")))
    if (len(vi_list) != len(gt_list)):
        raise ValueError("视频长度不匹配")
    total_psnr = 0
    for vi_name, gt_name in zip(vi_list, gt_list):
        vi, gt = get_image(vi_name, gt_name, crop=crop)
        total_psnr += base.psnr(vi, gt, mode=mode)

    return total_psnr / len(vi_list)


def get_focused_ssim(vi_path, gt_path, mode="RGB", crop=0):
    vi_list = sorted(glob.glob(os.path.join(vi_path, "*.png"))+glob.glob(os.path.join(vi_path, "*.jpg")))
    gt_list = sorted(glob.glob(os.path.join(gt_path, "*.png"))+glob.glob(os.path.join(gt_path, "*.jpg")))
    
    if (len(vi_list) != len(gt_list)):
        raise ValueError("视频长度不匹配")
    total_ssim = 0
    for vi_name, gt_name in zip(vi_list, gt_list):
        vi, gt = get_image(vi_name, gt_name, crop=crop)
        img_list, gt_list = get_focused_area_list(vi, gt)
        ssim_list = []
        for one_img, one_gt in zip(img_list, gt_list):
            ssim_list.append(base.ssim(one_img, one_gt, mode=mode))
        total_ssim += np.mean(ssim_list)

    return total_ssim / len(vi_list)


def get_ssim(vi_path, gt_path, mode="RGB", crop=0):
    vi_list = sorted(glob.glob(os.path.join(vi_path, "*.png")))
    gt_list = sorted(glob.glob(os.path.join(gt_path, "*.png")))
    if (len(vi_list) != len(gt_list)):
        raise ValueError("视频长度不匹配")
    total_ssim = 0
    for vi_name, gt_name in zip(vi_list, gt_list):
        vi, gt = get_image(vi_name, gt_name, crop=crop)
        total_ssim += base.ssim(vi, gt, mode=mode)

    return total_ssim / len(vi_list)


def get_ws_psnr(vi_path, gt_path, mode="RGB", crop=0):
    vi_list = sorted(glob.glob(os.path.join(vi_path, "*.png")))
    gt_list = sorted(glob.glob(os.path.join(gt_path, "*.png")))
    if (len(vi_list) != len(gt_list)):
        raise ValueError("视频长度不匹配")
    total_ws_psnr = 0

    for vi_name, gt_name in zip(vi_list, gt_list):
        vi, gt = get_image(vi_name, gt_name, crop=crop)
        total_ws_psnr += base.ws_psnr(vi, gt, mode=mode)

    return total_ws_psnr / len(vi_list)


def get_ws_ssim(vi_path, gt_path, mode="RGB", crop=0):
    vi_list = sorted(glob.glob(os.path.join(vi_path, "*.png")))
    gt_list = sorted(glob.glob(os.path.join(gt_path, "*.png")))
    if (len(vi_list) != len(gt_list)):
        raise ValueError("视频长度不匹配")
    total_ws_ssim = 0

    for vi_name, gt_name in zip(vi_list, gt_list):
        vi, gt = get_image(vi_name, gt_name, crop=crop)
        total_ws_ssim += base.ws_ssim(vi, gt, mode=mode)

    return total_ws_ssim / len(vi_list)
